#include "TreeClassifier.h"
//#include "SDLittleDog.h"
//#include "SDSlowWalk.h"
//#include "SDVectorOp.h"
//#include "Node.h"
#include <iostream>
#include <iomanip>
#include <ctime>
#include <stdlib.h>
#include <cstdlib>
#include <math.h>
#include <vector>
#include <fstream>
#include <stdio.h>
#include <utility>
#include <stack>
#include <queue>
#include <set>

using namespace std;

string TreeClassifier::classifierFileName = "TreeClassifier.txt";

TreeClassifier::TreeClassifier(string fileName, bool existingClassifier)
{
     correctlyLoaded = false;
     classifierRoot = NULL;
     if(!existingClassifier){
         cout << "CONSTRUCTOR: Generating classifier..." << endl;
         GenerateClassifier(fileName);
     }
     cout << "CONSTRUCTOR: Loading saved classifier..." << endl;
     LoadClassifier();
     cout << "CONSTRUCTOR: Testing classifier..." << endl;
     TestClassifier();
     cout << "CONSTRUCTOR FINISHED!" << endl;
}

float TreeClassifier::classify(vector<float>* costs)
{
     return FollowTree(costs, classifierRoot);
}

float TreeClassifier::FollowTree(vector<float>* costs, TreeNode* node)
{
     if(node->feature == -1) return node->p;
     else if(costs->at(node->feature) < node->tolerance) return FollowTree(costs, node->childL);
     else return FollowTree(costs, node->childR);
}

void TreeClassifier::GenerateClassifier(string fileName)
{
     //Open training file and read first line to obtain parameters
     cout << "LOADING DATA FILE '" << fileName << "'...";
     string line;
     ifstream trainingFile;
     trainingFile.open(fileName.c_str());
     if(!trainingFile.is_open()){
          printf("Error opening %s\n", fileName.c_str());
     }
     int spaceCount;
     if(getline(trainingFile,line)){
          int index =  line.find(" ");
          spaceCount = 0;
          while(index != string::npos){
                spaceCount++;
                index = line.find(" ",index+1);
          }
     } else {
          cout << fileName << " is an empty training file, quitting..." << endl;
          return;
     }
     cout << "number of features: " << spaceCount << endl;
     const int numFeatures = spaceCount;
     numberOfFeatures = numFeatures;

     //Declare data structures necessary to generate classifier
     set<vector<float>, smallest_f> data; //numFeatures + 1
/*     vector<vector<float> > trainingFeatures; //numFeatures
     vector<int> trainingOutcomes;
     vector<vector<float> > testingFeatures; //numFeatures
     vector<int> testingOutcomes;
     vector<vector<float> > tolerances; //dim [numFeatures][numTols] */

     //Read in data set
     //float firstDataPt[numFeatures + 2];
     vector<float> firstDataPt;
     int index1 = -1;  int index2 = line.find(" ");
     for(int i=0; i < numFeatures; i++){
           float feature = atof((line.substr(index1+1,index2-index1-1)).c_str());
//           cout << endl << (i+1) << "th feature is " << feature << endl;
           firstDataPt.push_back( feature );
           index1 = index2;
           index2 = line.find(" ",index1+1);
     }
     float outcome = atof((line.substr(index1)).c_str());
//     cout << "Outcome is " << outcome << endl << endl;
     firstDataPt.push_back( outcome );
     data.insert(firstDataPt);
     while(getline(trainingFile,line)){
          vector<float> dataPt;
          index1 = -1;  index2 = line.find(" ");
          for(int i=0; i < numFeatures; i++){
                float feature = atof((line.substr(index1+1,index2-index1-1)).c_str());
 //               cout << (i+1) << "th feature is " << feature << endl;
                dataPt.push_back( feature );
                index1 = index2;
                index2 = line.find(" ",index1+1);
          }
          outcome = atof((line.substr(index1)).c_str());
//          cout << "Outcome is " << outcome << endl;
          dataPt.push_back( outcome );
//          cout << dataPt.size() << endl;
          data.insert(dataPt);
//          cout << "data.size is " << data.size() << endl << endl;
//          getchar();
     }
     cout << "DONE" << endl;

     //Split data set into training and testing set
     cout << "ALOCATING TRAINING AND TEST DATA...";
     float normalizer = RAND_MAX;
     for(set<vector<float>, smallest_f>::iterator iter = data.begin(); iter != data.end(); iter++){
          int output;
          if((*iter).at(numFeatures) < .5) output = 0;
          else output = 1;
          vector<float> dataPt; //numFeatures
          for(int i=0; i < numFeatures; i++){
               dataPt.push_back( (*iter).at(i) );
          }
//          DisplayVector((*iter));
//          DisplayVector(dataPt);
//          cout << "OUTPUT: " << output << endl;
//          getchar();
          float randomNum = ((float) rand())/normalizer;
          if(randomNum < percentTraining){
               trainingFeatures.push_back(dataPt);
               trainingOutcomes.push_back(output);
          } else {
               testingFeatures.push_back(dataPt);
               testingOutcomes.push_back(output);
          }
//          cout << randomNum << " <? " << percentTraining << endl;
          cout << "Train: " << trainingFeatures.size() << ", " << trainingOutcomes.size() << endl;
          cout << "Test: " << testingFeatures.size() << ", " << testingOutcomes.size() << endl << endl;
     }
     cout << "DONE" << endl;

     //Determine tolerance values
     cout << "DETERMINING FEATURE TOLERANCES FROM DATA...";
     int groupSize = trainingFeatures.size()/(numTols + 1) + 1;
     for(int i=0; i < numFeatures; i++){
          multiset<float> featureSet;
          vector<float> featureTols;
          for(vector<vector<float> >::iterator iter = trainingFeatures.begin(); iter != trainingFeatures.end(); iter++){
               featureSet.insert((*iter).at(i));
          }
          float lowertol;
          float uppertol;
          multiset<float>::iterator iter = featureSet.begin();
          for(int j = 0; j < numTols; j++){
               for(int k = 0; k < groupSize -1; k++) iter++;
               lowertol = *iter;
               iter++;
               uppertol = *iter;
               featureTols.push_back( (lowertol+uppertol)/2 );
          }
//          DisplayVector(featureTols); getchar();
          tolerances.push_back(featureTols);
     }
     cout << "DONE" << endl;

     //Initialize first classifier node
     cout << "GENERATING CLASSIFIER...";
     int nodeIndexer = 0;
     TempTreeNode* firstNode = new TempTreeNode();
     for(int i=0; i < trainingFeatures.size(); i++){
          firstNode->dataFlow->push_back(true);
          if(trainingOutcomes.at(i) == 1) firstNode->good++;
          else firstNode->bad++;
     }
     firstNode->index = nodeIndexer;
     nodeIndexer++;
     //DisplayTempTreeNode(firstNode);
     //getchar();

     //Recursively generate classifier structure
     int nodeCount = Split(firstNode, &nodeIndexer);
     cout << "DONE" << endl;    

     //Save generated classifier
     cout << "SAVING CLASSIFIER TO FILE...";
     ofstream classifierFile(classifierFileName.c_str());
     if(!classifierFile){
        printf("Error opening the file!\n");
        return;
     } else {
          classifierFile<<setprecision(8);
          classifierFile<<fixed;
          classifierFile << nodeCount << endl;
          PrintNodesToFile(&classifierFile, firstNode);                
     }
     cout << "DONE" << endl;
     
/*
     FILE * classifierFile;
     classifierFile = fopen(classifierFileName.c_str(), "w");
     printf("Generating classifier and printing it to file \n");
     if(!classifierFile){
           cout << "Error opening the file!  " << classifierFileName << endl;
     } else {
          classifierFile << nodeCount << endl;
          int nodeCounter = 0;
          PrintNodesToFile(classifierFile, &firstNode, &nodeCounter);  
          cout << "FINISHED PRINTING CLASSIFIER TO FILE: " << (nodeCounter+1) << " nodes." << endl;                 
     }*/
}

void TreeClassifier::LoadClassifier()
{
     string line;
     ifstream classifier;
     classifier.open(classifierFileName.c_str());
     if(!classifier.is_open()){
          printf("Error opening %s\n", classifierFileName.c_str());
     }

     //declare parameters
     getline(classifier,line);
     const int numNodes = atoi(line.c_str());
     int children[numNodes][2];
     TreeNode* nodeAddresses[numNodes];
     
     //load nodes
     while(getline(classifier,line)){
          int index1 = line.find(" ");
          int index2 = line.find(" ", index1+1);
          int index3 = line.find(" ", index2+1);
          int index4 = line.find(" ", index3+1);
          int index5 = line.find(" ", index4+1);
          int nodeIndex = atoi((line.substr(0,index1)).c_str());
          TreeNode* node = new TreeNode();
          nodeAddresses[nodeIndex] = node;
          node->feature = atoi((line.substr(index1+1,index2-index1-1)).c_str());
          node->tolerance = atof((line.substr(index2+1,index3-index2-1)).c_str());
          children[nodeIndex][0] = atoi((line.substr(index3+1,index4-index3-1)).c_str());
          children[nodeIndex][1] = atoi((line.substr(index4+1,index5-index4-1)).c_str());
          node->p = atof((line.substr(index5)).c_str());
     }

     //Validate classifier structure
     if(/*validateClassifier(nodeAddresses, children, numNodes)*/ true){
          correctlyLoaded = true;
          //build actual node structure and store root node address
          for(int i = 0; i < numNodes; i++){
               if(children[i][0] == -1){
                    nodeAddresses[i]->childL = NULL;
                    nodeAddresses[i]->childR = NULL;
               } else {
                    nodeAddresses[i]->childL = nodeAddresses[children[i][0]];
                    nodeAddresses[i]->childR = nodeAddresses[children[i][1]];
               }
          }
          classifierRoot = nodeAddresses[0];
          cout << "Valid classifier loaded; has " << numNodes << " node(s)." << endl;
     } else {
          cout << "Classifier with " << numNodes << "node(s) is not valid" << endl;
     }
     classifier.close();
}

void TreeClassifier::TestClassifier(){
     cout << "TESTING THE CLASSIFIER..." << endl;
     int correct_stable = 0; //correctly predicted dog was stable
     int correct_flipped = 0; //correctly predicted dog flipped
     int wrong_stable = 0; //incorrectly predicted dog would be stable
     int wrong_flipped = 0; //incorrectly predicted dog would flip
     for(int i = 0; i < testingFeatures.size(); i++){
          float expectation = classify( &(testingFeatures.at(i)) );
          int prediction;
          if(expectation < .5) prediction = 0;
          else prediction = 1;
          if(prediction == 1){
               if(testingOutcomes.at(i) == 1) correct_stable++;
               else wrong_stable++;
          } else {
               if(testingOutcomes.at(i) == 1) wrong_flipped++;
               else correct_flipped++;
          }
     }
     float overall_correct = ((float) (correct_stable + correct_flipped))/((float) testingFeatures.size());
     cout << "     Overall correct: " << overall_correct << endl;
     cout << "     Successes predicted as successes/total: " << correct_stable << "/" << testingFeatures.size() << endl;
     cout << "     Successes predicted as flips/total: " << wrong_flipped << "/" << testingFeatures.size() << endl;
     cout << "     Flips predicted as successes/total: " << wrong_stable << "/" << testingFeatures.size() << endl;
     cout << "     Flips predicted as flips/total: " << correct_flipped << "/" << testingFeatures.size() << endl;
}

//This function determines whether or not to split the node passed to it
//and at what feature and tolerance to split at.
int TreeClassifier::Split(TempTreeNode* node, int* nodeIndexer)
{
     //DisplayTempTreeNode(node);
     float bestEntropy = h(node);
     cout << endl << endl << "---------Begin splitting of node " << node->index << " --------------" << endl << endl << "base entropy: " << bestEntropy << endl;
     int flowCount = node->good + node->bad;
     cout << "FLOW COUNT: " << flowCount << " = " << node->good << " good + " << node->bad << " bad" << endl;
     if(flowCount == 0){
          node->node.p = h_val(node->good, node->bad);
          node->left = NULL;
          node->right = NULL;
          //DisplayTempTreeNode(node);
          //getchar();
          return 1;
     }
     int bestFeature = -1;
     float bestTol = -1;
     int best_left_good = -1; int best_left_bad = -1;
     int best_right_good = -1; int best_right_bad = -1;
     vector<bool>* bestLeft = new vector<bool>();
     vector<bool>* bestRight = new vector<bool>();
     for(int i=0; i < numberOfFeatures; i++){
          for(int j=0; j < numTols; j++){
               float tol = tolerances.at(i).at(j);
               int left_good = 0; int left_bad = 0;
               int right_good = 0; int right_bad = 0;
               vector<bool>* leftVec = new vector<bool>();
               vector<bool>* rightVec = new vector<bool>();
               //DisplayTolerances();
               //cout << "tol at feature " << i << ", tol " << j << ": " << tol << endl;
               //getchar();
              // left_good = 0; left_bad = 0;
             //  right_good = 0; right_bad = 0;
               for(int k=0; k < trainingFeatures.size(); k++){
                    //cout << "data pt " << k << " has visibility " << node->dataFlow->at(k) << ": lg=" << left_good << ", lb=" << left_bad << ", rg=" << right_good << ", rb=" << right_bad << endl;
                    //cout << "data pt feature " << i << " value = " << trainingFeatures.at(k).at(i) << ", tol = " << tol << ", outcome = " << trainingOutcomes.at(k) << endl;
                    if(node->dataFlow->at(k)){
                         if(trainingFeatures.at(k).at(i) < tol){
                              //cout << endl << "Found that " << trainingFeatures.at(k).at(i) << " < " << tol << endl;
                              //getchar();
                              leftVec->push_back(true);
                              rightVec->push_back(false);
                              if(trainingOutcomes.at(k) == 1) left_good++; //cout << "left_good++" << endl;}
                              else left_bad++; //cout << "left_bad++" << endl;}
                         } else {
                              //cout << endl << "Found that " << trainingFeatures.at(k).at(i) << " => " << tol << endl;
                              leftVec->push_back(false);
                              rightVec->push_back(true);
                              if(trainingOutcomes.at(k) == 1) right_good++; //cout << "right_good++" << endl;}
                              else right_bad++; //cout << "right_bad++" << endl;}
                         }
                    } else {
                         leftVec->push_back(false);
                         rightVec->push_back(false);
                    }
                    //getchar();
               }
               float newEntropy = h_val(left_good, left_bad) + h_val(right_good,right_bad);
               
               if(newEntropy < bestEntropy){
                    cout << "FOUND BETTER ENTROPY!!!  at entropy " << newEntropy << endl;
                    cout << "feature= " << i << ", tol= " << j  << " (" << tol << "), lg=" << left_good << ", lb=" << left_bad << ", rg=" << right_good << ", rb=" << right_bad << ", vs. good=" << node->good << ", bad=" << node->bad << endl;
               cout << "newEntropy = " << newEntropy << ", oldEntropy = " << bestEntropy << endl;
                    bestFeature = i;
                    bestTol = tolerances.at(i).at(j);
                    bestEntropy = newEntropy;
                    best_left_good = left_good; best_left_bad = left_bad;
                    best_right_good = right_good; best_right_bad = right_bad;
                    delete bestLeft;
                    delete bestRight;
                    bestLeft = leftVec;
                    bestRight = rightVec;
               } else {
                    delete leftVec;
                    delete rightVec;
               }
               //cout << "bestFeature = " << bestFeature << ", bestTol = " << bestTol << ", bestEntropy = " << bestEntropy << endl << endl << endl << endl;
               //getchar();
          }
     }
     node->node.feature = bestFeature;
     node->node.tolerance = bestTol;

     //If not splitting on node, then finalize node fields
     if(bestFeature == -1){
          cout << "SENTENCE TO LEAF NODE: entropy = " << bestEntropy << endl;
          node->node.p = CalculateProbability(node->good, node->bad);
          //node->left = NULL;
          //node->right = NULL;
          //DisplayTempTreeNode(node);
          //getchar();
          DisplayTempTreeNode(node);
          cout << endl << "---------End splitting of node " << node->index << " --------------" << endl << endl;
          getchar();
          return 1;

     //If splitting on node, then finalize node fields, generate children and split children
     } else {
          cout << "SPLITTING!!! at feature " << bestFeature << " and tol " << bestTol << endl;
          //node->node.p = -1;
          TempTreeNode* newLeft = new TempTreeNode();
          TempTreeNode* newRight = new TempTreeNode();
          InitializeTempTreeNode(newLeft, nodeIndexer, best_left_good, best_left_bad, bestLeft);
          InitializeTempTreeNode(newRight, nodeIndexer, best_right_good, best_right_bad, bestRight);
          node->left = newLeft;
          node->right = newRight;
          DisplayTempTreeNode(node);
          DisplayTempTreeNode(newLeft);
          DisplayTempTreeNode(newRight);
          cout << endl << "---------End splitting of node " << node->index << " --------------" << endl << endl;
          getchar();
          return Split(newRight, nodeIndexer) + Split(newLeft, nodeIndexer) + 1;
     }
}

float TreeClassifier::h(TempTreeNode* node)
{
     return h_val(node->good, node->bad);
}

float TreeClassifier::h_val(int good, int bad)
{
     float prob_good = CalculateProbability(good, bad);
     float prob_bad = CalculateProbability(bad, good);
     float goodentropy = 0;
     float badentropy = 0;
     if(prob_good != 0) goodentropy = -prob_good*log10(prob_good);
     if(prob_bad != 0) badentropy = -prob_bad*log10(prob_bad);
     return  goodentropy+badentropy;
}

float TreeClassifier::CalculateProbability(int target, int other)
{
     int total = target+other;
     if(total == 0){
          target = 1;
          other = 1;
          total = 2;
     }
     return ((float) target)/((float) total);
}

bool TreeClassifier::is_loaded()
{
     return correctlyLoaded;
}

void TreeClassifier::PrintNodesToFile(ofstream* writer, TempTreeNode* node)
{
     int left = -1;
     int right = -1;
     if(node->left != NULL){
          left = node->left->index;
          right = node->right->index;
     }
     *writer << node->index << " " << node->node.feature << " "
               << node->node.tolerance << " " << left << " " << right
               << " " << node->node.p << endl;
     if(left != -1){
          PrintNodesToFile(writer, node->left);
          PrintNodesToFile(writer, node->right);
     }
}

int main(int argc, char** argv){
//DEBUGGING THE CLASSIFIER CODE//////////////////////////////////////////////
      cout << "Making the classifier..." << endl;
      TreeClassifier* classifier = new TreeClassifier("matlab/training.txt", false);
      cout << "Succeeded?  " << classifier->is_loaded() << endl << endl << endl;  

      int one1 = 0;
      int one2 = 1;
      int two1 = 0; int two2 = 0;
      int three1 = 1; int three2 = 8423817;
      cout << classifier->h_val(one1, one2) << endl;
      cout << classifier->h_val(two1, two2) << endl;
      cout << classifier->h_val(three1, three2) << endl;
}

void TreeClassifier::DisplayVector(vector<float> vect)
{
     cout << " vector size " << vect.size() << endl;
     for(int i = 0; i < vect.size(); i++){
          cout << vect.at(i) << " || ";
     } cout << endl;
}

void TreeClassifier::DisplayTempTreeNode(TempTreeNode* node)
{
     int left = 0;
     int right = 0;
     if(node->left == 0) left = -1;
     else left = node->left->index;
     if(node->right == NULL) right = -1;
     else right = node->right->index;

     cout << "index = " << node->index << endl;
     cout << "good = " << node->good << endl;
     cout << "bad = " << node->bad << endl;
     cout << "left = " << /*node->*/left << endl;
     cout << "right = " << /*node->*/right << endl;
     cout << "dataFlow->size() = " << node->dataFlow->size() << endl;
     DisplayTreeNode(&(node->node));
}

void TreeClassifier::DisplayTreeNode(TreeNode* node)
{
     cout << "node.feature = " << node->feature << endl;
     cout << "node.tolerance = " << node->tolerance << endl;
     cout << "node.p = " << node->p << endl;
     string left; string right;
     if(node->childL == NULL) left = right = "NULL";
     else left = right = "not null!";
     cout << "node.childL = " << left << endl;
     cout << "node.childR = " << right << endl << endl;
}

void TreeClassifier::DisplayTolerances()
{
     for(int i=0; i<tolerances.size(); i++){
          DisplayVector(tolerances.at(i));
     }
}

void TreeClassifier::InitializeTempTreeNode(TempTreeNode* node, int* indexer, int good, int bad, vector<bool>* flow)
{
     //cout << endl << endl << "NEW CHILD NODE:::::::::::::::::::::::::::::::::::::::" << endl << endl;
     node->index = *indexer;
     (*indexer)++;
     node->good = good;
     node->bad = bad;
     node->dataFlow = flow;



/*for(int i=0; i < trainingFeatures.size(); i++){ //iterate through data set
          if(parentDataFlow->at(i)){ //In flow?
               if(trainingFeatures.at(i).at(feature) < tolerance){ //going left?
cout << endl << "Found that " << trainingFeatures.at(i).at(feature) << " < " << tolerance << endl;
getchar();
                    if(which){ //if this is the left child node
                         node->dataFlow->push_back(true);
                         if(trainingOutcomes.at(i) == 1) node->good++;
                         else node->bad++;
                    } else { //this is the right node
                         node->dataFlow->push_back(false);
                    }
               } else { //going right
cout << endl << "Found that " << trainingFeatures.at(i).at(feature) << " => " << tolerance << endl;
                    if(which){ //this is left node
                         node->dataFlow->push_back(false);
                    } else { //this is the right node
                         node->dataFlow->push_back(true);
                         if(trainingOutcomes.at(i) == 1) node->good++;
                         else node->bad++;
                    }
               }
          } else { //out of flow
               node->dataFlow->push_back(false);
          }
     }
     //DisplayTempTreeNode(node);



 /*         int lef, rig; lef = rig = 0;
          for(int i=0; i < trainingFeatures.size(); i++){
               if(node->dataFlow->at(i)){
                    if(trainingFeatures.at(i).at(bestFeature) < bestTol){
                         newLeft->dataFlow->push_back(true);
                         newRight->dataFlow->push_back(false);
                         lef++;
                         if(trainingOutcomes.at(i) == 1) newLeft->good++;
                         else newLeft->bad++; 
                    } else {
                         newLeft->dataFlow->push_back(false);
                         newRight->dataFlow->push_back(true);
                         rig++;
                         if(trainingOutcomes.at(i) == 1) newRight->good++;
                         else newRight->bad++; 
                    }
               } else {
                    newLeft->dataFlow->push_back(false);
                    newRight->dataFlow->push_back(false);
               }
          }
          cout << "Varify correct flow: " << lef << " left = " << newLeft->good << " good + " << newLeft->bad << " bad." << endl;
          cout << "Varify correct flow: " << rig << " right = " << newRight->good << " good + " << newRight->bad << " bad." << endl;
          (*nodeIndexer)++;
          newLeft->index = *nodeIndexer;
          (*nodeIndexer)++;
          newRight->index = *nodeIndexer;





node = new TempTreeNode();
     node->index = *indexer;
     cout << endl << endl << "address: " << indexer << ", value: " << *indexer << endl;
     (*indexer)++;
     node->good = 0;
     node->bad = 0;
     node->left = NULL; cout << "Made left null: " << node->left << endl;
     node->right = NULL;
     node->dataFlow = new vector<bool>();
     cout << "just initialized dataFlow vector" << endl;
     cout << "dataFLow.size() = " << node->dataFlow->size() << endl;
     cout << "just displayed data flow size" << endl;
     node->node.feature = -1;
     node->node.tolerance = -1;
     node->node.p = -1;
     node->node.childL = NULL;
     node->node.childR = NULL;*/
}
